package business.server;

import demo.storage.LogEntry;

import java.io.*;
import java.security.MessageDigest;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 键值存储状态机实现
 * 
 * <p>实现了一个分布式键值存储，支持基本的CRUD操作：</p>
 * <ul>
 *   <li><b>SET key value</b>：设置键值对</li>
 *   <li><b>GET key</b>：获取键对应的值</li>
 *   <li><b>DELETE key</b>：删除键值对</li>
 *   <li><b>EXISTS key</b>：检查键是否存在</li>
 *   <li><b>CLEAR</b>：清空所有数据</li>
 * </ul>
 * 
 * <p><b>典型业务场景：</b></p>
 * <ul>
 *   <li>分布式配置中心</li>
 *   <li>分布式缓存</li>
 *   <li>服务注册与发现</li>
 *   <li>分布式锁服务</li>
 * </ul>
 * 
 * @author Raft Implementation Team
 * @version 1.0
 * @since 1.0
 */
public class KeyValueStateMachine implements StateMachine {
    
    /** 存储键值对的并发安全Map */
    private final Map<String, String> store = new ConcurrentHashMap<>();
    
    /** 已应用的最后一个日志索引 */
    private volatile int lastAppliedIndex = 0;
    
    /** 命令解析的正则表达式 */
    private static final Pattern SET_PATTERN = Pattern.compile("SET\\s+(\\S+)\\s+(.+)", Pattern.CASE_INSENSITIVE);
    private static final Pattern GET_PATTERN = Pattern.compile("GET\\s+(\\S+)", Pattern.CASE_INSENSITIVE);
    private static final Pattern DELETE_PATTERN = Pattern.compile("DELETE\\s+(\\S+)", Pattern.CASE_INSENSITIVE);
    private static final Pattern EXISTS_PATTERN = Pattern.compile("EXISTS\\s+(\\S+)", Pattern.CASE_INSENSITIVE);
    private static final Pattern CLEAR_PATTERN = Pattern.compile("CLEAR", Pattern.CASE_INSENSITIVE);
    
    @Override
    public Object apply(LogEntry entry) throws StateMachineException {
        if (entry == null || entry.getCommand().trim().isEmpty()) {
            return "OK"; // 空命令，忽略
        }
        
        String command = entry.getCommand().trim();
        Object result;
        
        try {
            // 解析并执行命令
            if (SET_PATTERN.matcher(command).matches()) {
                result = handleSet(command);
            } else if (GET_PATTERN.matcher(command).matches()) {
                result = handleGet(command);
            } else if (DELETE_PATTERN.matcher(command).matches()) {
                result = handleDelete(command);
            } else if (EXISTS_PATTERN.matcher(command).matches()) {
                result = handleExists(command);
            } else if (CLEAR_PATTERN.matcher(command).matches()) {
                result = handleClear();
            } else {
                throw new StateMachineException("未知命令: " + command);
            }
            
            // 更新已应用索引
            lastAppliedIndex = entry.getIndex();
            
            System.out.println("状态机应用命令: " + command + " -> " + result);
            return result;
            
        } catch (Exception e) {
            throw new StateMachineException("执行命令失败: " + command, e);
        }
    }
    
    /**
     * 处理SET命令
     */
    private String handleSet(String command) {
        Matcher matcher = SET_PATTERN.matcher(command);
        if (matcher.find()) {
            String key = matcher.group(1);
            String value = matcher.group(2);
            store.put(key, value);
            return "OK";
        }
        return "ERROR: Invalid SET command format";
    }
    
    /**
     * 处理GET命令
     */
    private String handleGet(String command) {
        Matcher matcher = GET_PATTERN.matcher(command);
        if (matcher.find()) {
            String key = matcher.group(1);
            String value = store.get(key);
            return value != null ? value : "NULL";
        }
        return "ERROR: Invalid GET command format";
    }
    
    /**
     * 处理DELETE命令
     */
    private String handleDelete(String command) {
        Matcher matcher = DELETE_PATTERN.matcher(command);
        if (matcher.find()) {
            String key = matcher.group(1);
            String removedValue = store.remove(key);
            return removedValue != null ? "OK" : "NOT_FOUND";
        }
        return "ERROR: Invalid DELETE command format";
    }
    
    /**
     * 处理EXISTS命令
     */
    private String handleExists(String command) {
        Matcher matcher = EXISTS_PATTERN.matcher(command);
        if (matcher.find()) {
            String key = matcher.group(1);
            return store.containsKey(key) ? "TRUE" : "FALSE";
        }
        return "ERROR: Invalid EXISTS command format";
    }
    
    /**
     * 处理CLEAR命令
     */
    private String handleClear() {
        int count = store.size();
        store.clear();
        return "OK (" + count + " keys removed)";
    }
    
    @Override
    public Snapshot createSnapshot(int lastIncludedIndex, int lastIncludedTerm) {
        try {
            // 序列化当前状态
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            ObjectOutputStream oos = new ObjectOutputStream(baos);
            
            // 写入元数据
            oos.writeInt(lastIncludedIndex);
            oos.writeInt(lastIncludedTerm);
            oos.writeInt(lastAppliedIndex);
            
            // 写入存储数据
            oos.writeInt(store.size());
            for (Map.Entry<String, String> entry : store.entrySet()) {
                oos.writeUTF(entry.getKey());
                oos.writeUTF(entry.getValue());
            }
            
            oos.close();
            byte[] data = baos.toByteArray();
            
            // 计算校验和
            String checksum = calculateChecksum(data);
            
            return new KeyValueSnapshot(lastIncludedIndex, lastIncludedTerm, data, checksum);
            
        } catch (IOException e) {
            throw new RuntimeException("创建快照失败", e);
        }
    }
    
    @Override
    public void restoreFromSnapshot(Snapshot snapshot) {
        try {
            byte[] data = snapshot.getData();
            
            // 验证校验和
            String expectedChecksum = snapshot.getChecksum();
            String actualChecksum = calculateChecksum(data);
            if (!expectedChecksum.equals(actualChecksum)) {
                throw new RuntimeException("快照数据校验失败");
            }
            
            // 反序列化状态
            ByteArrayInputStream bais = new ByteArrayInputStream(data);
            ObjectInputStream ois = new ObjectInputStream(bais);
            
            // 读取元数据
            int lastIncludedIndex = ois.readInt();
            int lastIncludedTerm = ois.readInt();
            int restoredLastAppliedIndex = ois.readInt();
            
            // 读取存储数据
            store.clear();
            int size = ois.readInt();
            for (int i = 0; i < size; i++) {
                String key = ois.readUTF();
                String value = ois.readUTF();
                store.put(key, value);
            }
            
            // 更新状态
            this.lastAppliedIndex = restoredLastAppliedIndex;
            
            ois.close();
            
            System.out.println("从快照恢复状态机: " + size + " 个键值对, lastAppliedIndex=" + lastAppliedIndex);
            
        } catch (IOException e) {
            throw new RuntimeException("从快照恢复失败", e);
        }
    }
    
    @Override
    public Map<String, Object> getState() {
        Map<String, Object> state = new HashMap<>();
        state.put("storeSize", store.size());
        state.put("lastAppliedIndex", lastAppliedIndex);
        state.put("memoryUsage", Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory());
        
        // 添加前10个键值对用于调试
        List<String> sampleKeys = store.keySet().stream()
                .limit(10)
                .collect(java.util.stream.Collectors.toList());
        state.put("sampleKeys", sampleKeys);
        
        return state;
    }
    
    @Override
    public boolean isValidCommand(String command) {
        if (command == null || command.trim().isEmpty()) {
            return false;
        }
        
        String trimmed = command.trim();
        return SET_PATTERN.matcher(trimmed).matches() ||
               GET_PATTERN.matcher(trimmed).matches() ||
               DELETE_PATTERN.matcher(trimmed).matches() ||
               EXISTS_PATTERN.matcher(trimmed).matches() ||
               CLEAR_PATTERN.matcher(trimmed).matches();
    }
    
    @Override
    public int getLastAppliedIndex() {
        return lastAppliedIndex;
    }
    
    /**
     * 获取存储的键值对数量
     */
    public int size() {
        return store.size();
    }
    
    /**
     * 获取所有键
     */
    public Set<String> keySet() {
        return new HashSet<>(store.keySet());
    }
    
    /**
     * 直接获取值（用于读操作，不经过Raft日志）
     */
    public String getValue(String key) {
        return store.get(key);
    }
    
    /**
     * 计算数据的MD5校验和
     */
    private String calculateChecksum(byte[] data) {
        try {
            MessageDigest md = MessageDigest.getInstance("MD5");
            byte[] hash = md.digest(data);
            StringBuilder sb = new StringBuilder();
            for (byte b : hash) {
                sb.append(String.format("%02x", b));
            }
            return sb.toString();
        } catch (Exception e) {
            throw new RuntimeException("计算校验和失败", e);
        }
    }
    
    /**
     * 键值存储快照实现
     */
    private static class KeyValueSnapshot implements Snapshot {
        private final int lastIncludedIndex;
        private final int lastIncludedTerm;
        private final byte[] data;
        private final String checksum;
        private final long timestamp;
        
        public KeyValueSnapshot(int lastIncludedIndex, int lastIncludedTerm, byte[] data, String checksum) {
            this.lastIncludedIndex = lastIncludedIndex;
            this.lastIncludedTerm = lastIncludedTerm;
            this.data = data.clone(); // 防御性复制
            this.checksum = checksum;
            this.timestamp = System.currentTimeMillis();
        }
        
        @Override
        public int getLastIncludedIndex() {
            return lastIncludedIndex;
        }
        
        @Override
        public int getLastIncludedTerm() {
            return lastIncludedTerm;
        }
        
        @Override
        public byte[] getData() {
            return data.clone(); // 防御性复制
        }
        
        @Override
        public String getChecksum() {
            return checksum;
        }
        
        @Override
        public long getTimestamp() {
            return timestamp;
        }
    }
} 